#include <cassert>
#include <bitset>
#include <string>

#include <openssl/bn.h>
#include <openssl/ec.h>
#include <glog/logging.h>

std::string bn2str(BIGNUM* bn_x) {
  int size = BN_num_bytes(bn_x);
  return std::string(64 - 2*size, '0') + BN_bn2hex(bn_x);
}

bool VerifyEcdsa0(
    const EC_GROUP* group,
    const EC_POINT* public_key,
    const EC_POINT* generator,
    const BIGNUM* z,
    const BIGNUM* curve_order,
    const BIGNUM* upper_bound,
    const std::pair<BIGNUM*, BIGNUM*>& sig) {
  const BIGNUM* r = sig.first;
  const BIGNUM* w = sig.second;
  // z, r, w should be smaller than 2^251.
  assert(!BN_is_zero(z) && "Message cannot be zero.");
  assert(BN_cmp(z, upper_bound) == -1 && "z is too big.");
  assert(!BN_is_zero(r) && "r cannot be zero.");
  assert(BN_cmp(r, upper_bound) == -1 && "r is too big.");
  assert(!BN_is_zero(w) && "w cannot be zero.");
  assert(BN_cmp(w, upper_bound) == -1 && "w is too big.");

  BIGNUM* zw = BN_new();
  BIGNUM* rw = BN_new();
  BIGNUM* x1 = BN_new();
  BIGNUM* x2 = BN_new();
  EC_POINT* zw_g = EC_POINT_new(group);
  EC_POINT* rw_q = EC_POINT_new(group);
  EC_POINT* tmp = EC_POINT_new(group);
  BN_CTX* bn_ctx = BN_CTX_new();

  BN_mod_mul(rw, r, w, curve_order, bn_ctx);
  BN_mod_mul(zw, z, w, curve_order, bn_ctx);
  EC_POINT_mul(group, zw_g, NULL, generator, zw, bn_ctx);
  EC_POINT_mul(group, rw_q, NULL, public_key, rw, bn_ctx);

  // (zw_g + rw_q).x
  EC_POINT_add(group, tmp, zw_g, rw_q, bn_ctx);
  EC_POINT_get_affine_coordinates(group, tmp, x1, NULL, bn_ctx);

  // (zw_g - rw_q).x
  EC_POINT_invert(group, rw_q, bn_ctx); // -rw_q
  EC_POINT_add(group, tmp, zw_g, rw_q, bn_ctx);
  EC_POINT_get_affine_coordinates(group, tmp, x2, NULL, bn_ctx);

  bool res = !(BN_cmp(x1, r) & BN_cmp(x2, r));

  BN_free(x1);
  BN_free(x2);
  BN_free(zw);
  BN_free(rw);

  EC_POINT_free(zw_g);
  EC_POINT_free(rw_q);
  EC_POINT_free(tmp);
  BN_CTX_free(bn_ctx);
  return res;
}

std::string SignEcdsa0(
    const EC_GROUP* group,
    const EC_POINT* public_key,
    const EC_POINT* generator,
    const BIGNUM* privkey,
    const BIGNUM* z,
    const BIGNUM* curve_order,
    const BIGNUM* upper_bound,
    const bool verify) {
  EC_POINT* kG = EC_POINT_new(group);
  BIGNUM* randk = BN_new();
  BIGNUM* r = BN_new();
  BIGNUM* s = BN_new();
  BIGNUM* w = BN_new();
  BN_CTX* bn_ctx = BN_CTX_new();

  while (true) {
    BN_rand_range(randk, curve_order);
    // DLOG(INFO) << "randk: " << bn2str(randk);

    EC_POINT_mul(group, kG, NULL, generator, randk, bn_ctx);
    EC_POINT_get_affine_coordinates(group, kG, r, NULL, bn_ctx);
    if (!((BN_cmp(r, curve_order) == -1) && (!BN_is_zero(r)))) {
      // DLOG(WARNING) << "Bad randomness, please try a different a different k.";
      continue;
    }
    // DLOG(INFO) << "r: " << bn2str(r);

    BN_mod_mul(s, r, privkey, curve_order, bn_ctx);
    BN_mod_add_quick(s, s, z, curve_order);
    // DLOG(INFO) << "s: " << bn2str(s);

    BN_mod_inverse(w, s, curve_order, bn_ctx);
    BN_mod_mul(w, randk, w, curve_order, bn_ctx);
    if (BN_cmp(w, upper_bound) != -1) {
      // DLOG(WARNING) << "Bad randomness, please try a different k.";
      continue;
    }
    // DLOG(INFO) << "w: " << bn2str(w);

    BN_mod_inverse(s, w, curve_order, bn_ctx);
    // DLOG(INFO) << "s: " << bn2str(s);

    auto res = bn2str(r) + bn2str(s);
    // DLOG(INFO) << "res: " << res;

    if (verify) {
      bool vres = VerifyEcdsa0(group, public_key, generator, z, curve_order, upper_bound, {r, w});
      if (!vres) {
        LOG(INFO) << "wrong";
        continue;
      }
    }

    BN_free(s);
    BN_free(w);
    BN_free(r);
    BN_free(randk);

    EC_POINT_free(kG);
    BN_CTX_free(bn_ctx);
    return res;
  }
}
